import type { Plugin } from 'vite';
import { parse } from '@babel/parser';
import * as traverse from '@babel/traverse';
import generate from '@babel/generator';
import * as t from '@babel/types';

export default function removeConsoleSafely(): Plugin {
  return {
    name: 'vite-plugin-remove-console-safely',
    apply: 'build',
    transform(code, id) {
      if (!/\.(js|ts|jsx|tsx)$/.test(id)) {
        return;
      }
      const ast = parse(code, {
        sourceType: 'module',
        plugins: ['typescript', 'jsx']
      });

      traverse.default(ast, {
        ExpressionStatement(path) {
          const expr = path.node.expression;
          if (
            t.isCallExpression(expr) &&
            t.isMemberExpression(expr.callee) &&
            t.isIdentifier(expr.callee.object, { name: 'console' }) &&
            t.isIdentifier(expr.callee.property, { name: 'log' })
          ) {
            const parent = path.parentPath;
            if (
              t.isBlockStatement(parent.node) &&
              t.isArrowFunctionExpression(parent.parent) &&
              parent.node.body.length === 1
            ) {
              // 箭头函数体唯一 console.log 不删除
            } else {
              path.remove();
            }
          }
        }
      });

      const output = generate(ast, {}, code);

      return {
        code: output.code,
        map: null
      };
    }
  };
}
